"""
Hybrid Configuration Manager.

This module provides XML configuration parsing and validation for the
Hybrid PETSc-Kokkos extension, handling solver parameters and execution
settings compatible with the hybrid framework.
"""

import xml.etree.ElementTree as ET
import json
import logging
from pathlib import Path
from typing import Dict, List, Optional, Any, Union
from dataclasses import dataclass, field

from gcr_solver_manager.extensions.base_builder import BuildConfig
from gcr_solver_manager.extensions.base_runner import RunConfig


@dataclass
class HybridSolverConfig:
    """Configuration for hybrid solver execution."""
    
    # Solver parameters
    solver_type: str = "gcr"
    max_iterations: int = 100
    tolerance: float = 1e-6
    restart: Optional[int] = None
    
    # Problem configuration
    problem_type: str = "helmholtz"
    resolution: float = 1.0
    domain_size: List[float] = field(default_factory=lambda: [1.0, 1.0, 1.0])
    
    # Parallel configuration
    mpi_processes: int = 1
    gpu_backend: Optional[str] = None
    
    # Performance options
    debug_mode: bool = False
    verbose: bool = False
    benchmark_mode: bool = False
    
    # I/O configuration
    output_frequency: int = 10
    save_residual_history: bool = True
    export_solution: bool = False


@dataclass
class HybridBuildConfig:
    """Configuration for hybrid framework builds."""
    
    # Build options
    build_type: str = "Release"
    enable_gpu: bool = True
    enable_mpi: bool = True
    enable_testing: bool = True
    enable_benchmarks: bool = False
    
    # Dependencies
    petsc_dir: Optional[str] = None
    petsc_arch: Optional[str] = None
    kokkos_root: Optional[str] = None
    
    # Compiler settings
    parallel_jobs: int = 4
    compiler_flags: Dict[str, str] = field(default_factory=dict)
    
    # Advanced options
    custom_cmake_args: List[str] = field(default_factory=list)


class HybridConfigManager:
    """Manager for hybrid framework configuration."""
    
    def __init__(self):
        """Initialize the configuration manager."""
        self.logger = logging.getLogger(__name__)
        self._schema_cache = {}
        
        # Define supported solvers and their parameters
        self.supported_solvers = {
            "gcr": {
                "description": "Generalized Conjugate Residual",
                "parameters": ["tolerance", "max_iterations", "restart"]
            },
            "ca-gcr": {
                "description": "Communication-Avoiding GCR",
                "parameters": ["tolerance", "max_iterations", "restart", "s_value"]
            },
            "gmres": {
                "description": "Generalized Minimal Residual",
                "parameters": ["tolerance", "max_iterations", "restart"]
            },
            "bca-gmres": {
                "description": "Block Communication-Avoiding GMRES",
                "parameters": ["tolerance", "max_iterations", "restart", "block_size"]
            }
        }
    
    def parse_xml_config(self, config_path: Path) -> HybridSolverConfig:
        """Parse XML configuration file."""
        self.logger.info(f"Parsing hybrid XML configuration: {config_path}")
        
        if not config_path.exists():
            raise FileNotFoundError(f"Configuration file not found: {config_path}")
        
        try:
            tree = ET.parse(config_path)
            root = tree.getroot()
            
            config = HybridSolverConfig()
            
            # Parse solver configuration
            self._parse_solver_section(root, config)
            
            # Parse problem configuration
            self._parse_problem_section(root, config)
            
            # Parse parallel configuration
            self._parse_parallel_section(root, config)
            
            # Parse performance options
            self._parse_performance_section(root, config)
            
            # Parse I/O configuration
            self._parse_io_section(root, config)
            
            self.logger.info(f"Parsed configuration: solver={config.solver_type}, "
                           f"iterations={config.max_iterations}, tolerance={config.tolerance}")
            
            return config
            
        except ET.ParseError as e:
            raise ValueError(f"Invalid XML configuration: {e}")
        except Exception as e:
            raise ValueError(f"Configuration parsing error: {e}")
    
    def _parse_solver_section(self, root: ET.Element, config: HybridSolverConfig):
        """Parse solver-specific configuration."""
        
        # Solver type
        solver_elem = root.find('solver_type')
        if solver_elem is not None:
            config.solver_type = solver_elem.text.strip()
        elif root.find('solver') is not None:
            config.solver_type = root.find('solver').text.strip()
        
        # Validate solver type
        if config.solver_type not in self.supported_solvers:
            raise ValueError(f"Unsupported solver type: {config.solver_type}")
        
        # Solver parameters
        solver_params = root.find('solver_parameters') or root.find('parameters')
        if solver_params is not None:
            
            # Max iterations
            max_iter = solver_params.find('max_iterations')
            if max_iter is not None:
                config.max_iterations = int(max_iter.text.strip())
            
            # Tolerance
            tolerance = solver_params.find('tolerance')
            if tolerance is not None:
                config.tolerance = float(tolerance.text.strip())
            
            # Restart (for GMRES-family solvers)
            restart = solver_params.find('restart')
            if restart is not None:
                config.restart = int(restart.text.strip())
    
    def _parse_problem_section(self, root: ET.Element, config: HybridSolverConfig):
        """Parse problem configuration."""
        problem_elem = root.find('problem')
        if problem_elem is not None:
            
            # Problem type
            prob_type = problem_elem.find('type') or problem_elem.find('problem_type')
            if prob_type is not None:
                config.problem_type = prob_type.text.strip()
            
            # Resolution
            resolution = problem_elem.find('resolution')
            if resolution is not None:
                config.resolution = float(resolution.text.strip())
            
            # Domain size
            domain = problem_elem.find('domain')
            if domain is not None:
                size_elem = domain.find('size')
                if size_elem is not None:
                    size_text = size_elem.text.strip()
                    if ',' in size_text:
                        config.domain_size = [float(x.strip()) for x in size_text.split(',')]
                    else:
                        size = float(size_text)
                        config.domain_size = [size, size, size]
    
    def _parse_parallel_section(self, root: ET.Element, config: HybridSolverConfig):
        """Parse parallel execution configuration."""
        parallel_elem = root.find('parallel')
        if parallel_elem is not None:
            
            # MPI processes
            mpi_procs = parallel_elem.find('mpi_processes') or parallel_elem.find('processes')
            if mpi_procs is not None:
                config.mpi_processes = int(mpi_procs.text.strip())
            
            # GPU backend
            gpu_backend = parallel_elem.find('gpu_backend') or parallel_elem.find('gpu')
            if gpu_backend is not None:
                backend_text = gpu_backend.text.strip().lower()
                if backend_text in ['cuda', 'hip', 'none', 'off']:
                    config.gpu_backend = None if backend_text in ['none', 'off'] else backend_text
    
    def _parse_performance_section(self, root: ET.Element, config: HybridSolverConfig):
        """Parse performance and debugging options."""
        
        # Debug mode
        debug_elem = root.find('debug') or root.find('debug_mode')
        if debug_elem is not None:
            config.debug_mode = self._parse_bool(debug_elem.text)
        
        # Verbose mode
        verbose_elem = root.find('verbose')
        if verbose_elem is not None:
            config.verbose = self._parse_bool(verbose_elem.text)
        
        # Benchmark mode
        benchmark_elem = root.find('benchmark') or root.find('benchmark_mode')
        if benchmark_elem is not None:
            config.benchmark_mode = self._parse_bool(benchmark_elem.text)
        
        # Performance section
        perf_elem = root.find('performance')
        if perf_elem is not None:
            debug = perf_elem.find('debug')
            if debug is not None:
                config.debug_mode = self._parse_bool(debug.text)
            
            verbose = perf_elem.find('verbose')
            if verbose is not None:
                config.verbose = self._parse_bool(verbose.text)
    
    def _parse_io_section(self, root: ET.Element, config: HybridSolverConfig):
        """Parse I/O configuration."""
        io_elem = root.find('io') or root.find('output')
        if io_elem is not None:
            
            # Output frequency
            freq_elem = io_elem.find('frequency') or io_elem.find('output_frequency')
            if freq_elem is not None:
                config.output_frequency = int(freq_elem.text.strip())
            
            # Save residual history
            residual_elem = io_elem.find('save_residuals') or io_elem.find('residual_history')
            if residual_elem is not None:
                config.save_residual_history = self._parse_bool(residual_elem.text)
            
            # Export solution
            export_elem = io_elem.find('export_solution') or io_elem.find('save_solution')
            if export_elem is not None:
                config.export_solution = self._parse_bool(export_elem.text)
    
    def _parse_bool(self, text: str) -> bool:
        """Parse boolean value from string."""
        if not text:
            return False
        
        text = text.strip().lower()
        return text in ['true', 'yes', '1', 'on', 'enabled']
    
    def convert_to_run_config(self, hybrid_config: HybridSolverConfig,
                             binary_path: Path,
                             config_file: Optional[Path] = None) -> RunConfig:
        """Convert hybrid configuration to GSM RunConfig."""
        
        run_config = RunConfig(
            binary_path=binary_path,
            config_file=config_file,
            solver_type=hybrid_config.solver_type,
            max_iterations=hybrid_config.max_iterations,
            tolerance=hybrid_config.tolerance,
            mpi_processes=hybrid_config.mpi_processes,
            gpu_backend=hybrid_config.gpu_backend,
            debug_mode=hybrid_config.debug_mode,
            verbose=hybrid_config.verbose
        )
        
        # Add extra arguments based on hybrid config
        extra_args = []
        
        if hybrid_config.restart:
            extra_args.extend(["--restart", str(hybrid_config.restart)])
        
        if hybrid_config.problem_type != "helmholtz":
            extra_args.extend(["--problem-type", hybrid_config.problem_type])
        
        if hybrid_config.resolution != 1.0:
            extra_args.extend(["--resolution", str(hybrid_config.resolution)])
        
        if hybrid_config.benchmark_mode:
            extra_args.append("--benchmark")
        
        if hybrid_config.output_frequency != 10:
            extra_args.extend(["--output-frequency", str(hybrid_config.output_frequency)])
        
        if not hybrid_config.save_residual_history:
            extra_args.append("--no-residual-history")
        
        if hybrid_config.export_solution:
            extra_args.append("--export-solution")
        
        run_config.extra_args = extra_args
        
        return run_config
    
    def convert_to_build_config(self, hybrid_build_config: HybridBuildConfig) -> BuildConfig:
        """Convert hybrid build configuration to GSM BuildConfig."""
        
        build_config = BuildConfig(
            debug_mode=(hybrid_build_config.build_type == "Debug"),
            parallel_jobs=hybrid_build_config.parallel_jobs,
            gpu_backend="cuda" if hybrid_build_config.enable_gpu else None,
            enable_mpi=hybrid_build_config.enable_mpi,
            enable_testing=hybrid_build_config.enable_testing
        )
        
        # Setup environment
        environment = {}
        
        if hybrid_build_config.petsc_dir:
            environment["PETSC_DIR"] = hybrid_build_config.petsc_dir
        
        if hybrid_build_config.petsc_arch:
            environment["PETSC_ARCH"] = hybrid_build_config.petsc_arch
        
        if hybrid_build_config.kokkos_root:
            environment["Kokkos_ROOT"] = hybrid_build_config.kokkos_root
        
        build_config.environment = environment
        
        # Add custom CMake arguments
        build_config.extra_cmake_args = hybrid_build_config.custom_cmake_args.copy()
        
        # Add compiler flags as CMake arguments
        for flag_type, flags in hybrid_build_config.compiler_flags.items():
            if flag_type == "cxx_flags":
                build_config.extra_cmake_args.append(f"-DCMAKE_CXX_FLAGS={flags}")
            elif flag_type == "c_flags":
                build_config.extra_cmake_args.append(f"-DCMAKE_C_FLAGS={flags}")
            elif flag_type == "cuda_flags":
                build_config.extra_cmake_args.append(f"-DCMAKE_CUDA_FLAGS={flags}")
        
        return build_config
    
    def validate_config(self, config: HybridSolverConfig) -> List[str]:
        """Validate hybrid solver configuration."""
        errors = []
        
        # Validate solver type
        if config.solver_type not in self.supported_solvers:
            errors.append(f"Unsupported solver type: {config.solver_type}")
        
        # Validate numeric parameters
        if config.max_iterations <= 0:
            errors.append("max_iterations must be positive")
        
        if config.tolerance <= 0:
            errors.append("tolerance must be positive")
        
        if config.mpi_processes <= 0:
            errors.append("mpi_processes must be positive")
        
        # Validate solver-specific parameters
        solver_info = self.supported_solvers.get(config.solver_type, {})
        if config.restart and "restart" not in solver_info.get("parameters", []):
            errors.append(f"restart parameter not supported for {config.solver_type}")
        
        # Validate GPU backend
        if config.gpu_backend and config.gpu_backend not in ["cuda", "hip"]:
            errors.append(f"Unsupported GPU backend: {config.gpu_backend}")
        
        return errors
    
    def can_handle_config(self, config_path: Path) -> bool:
        """Check if this manager can handle the given configuration file."""
        if not config_path.exists():
            return False
        
        # Check file extension
        if config_path.suffix.lower() not in ['.xml']:
            return False
        
        try:
            # Try to parse as XML and look for hybrid-specific elements
            tree = ET.parse(config_path)
            root = tree.getroot()
            
            # Look for hybrid/solver indicators
            indicators = [
                'solver_type', 'solver', 'solver_parameters',
                'problem', 'parallel', 'hybrid'
            ]
            
            for indicator in indicators:
                if root.find(indicator) is not None:
                    return True
            
            # Check if solver type is one we support
            solver_elem = root.find('solver_type') or root.find('solver')
            if solver_elem is not None:
                solver_type = solver_elem.text.strip()
                return solver_type in self.supported_solvers
            
            return False
            
        except (ET.ParseError, Exception):
            return False
    
    def create_template_config(self, solver_type: str = "gcr") -> str:
        """Create a template XML configuration for the specified solver."""
        template = f'''<?xml version="1.0" encoding="UTF-8"?>
<solver_config>
  <!-- Solver Configuration -->
  <solver_type>{solver_type}</solver_type>
  <solver_parameters>
    <max_iterations>100</max_iterations>
    <tolerance>1e-6</tolerance>
    {f'<restart>30</restart>' if solver_type in ['gmres', 'bca-gmres'] else ''}
  </solver_parameters>
  
  <!-- Problem Configuration -->
  <problem>
    <type>helmholtz</type>
    <resolution>1.0</resolution>
    <domain>
      <size>1.0, 1.0, 1.0</size>
    </domain>
  </problem>
  
  <!-- Parallel Configuration -->
  <parallel>
    <mpi_processes>1</mpi_processes>
    <gpu_backend>cuda</gpu_backend>
  </parallel>
  
  <!-- Performance Options -->
  <performance>
    <debug>false</debug>
    <verbose>false</verbose>
    <benchmark>false</benchmark>
  </performance>
  
  <!-- I/O Configuration -->
  <io>
    <output_frequency>10</output_frequency>
    <save_residuals>true</save_residuals>
    <export_solution>false</export_solution>
  </io>
</solver_config>'''
        
        return template